#include "processdevice.h"
#include "misc.h"
#include "logger.h"

#include <Tlhelp32.h>
#include "error.h"

const char *windowThumbnailClassName = "ThumbnailClass";



ProcessDevice::ProcessDevice()
{
    attachedHwnd = NULL;
    attached = false;
}

ProcessDevice::~ProcessDevice()
{
    if( attachedHwnd )
        detach();
    attached = false;
}

int ProcessDevice::attach(HWND hwnd)
{
    if( attachedHwnd )
        detach();

    bool success = AttachThreadInput(GetCurrentThreadId(),
        GetWindowThreadProcessId(hwnd, NULL), true);

    if( success ) {
        attachedHwnd = hwnd;
        attached = true;
    }

    return success;
}

int ProcessDevice::detach()
{
    if( !attachedHwnd )
        return false;

    bool success = AttachThreadInput(GetCurrentThreadId(),
        GetWindowThreadProcessId(attachedHwnd, NULL), false);

    attachedHwnd = NULL;
    attached = false;

    return success;
}

int ProcessDevice::reattach()
{
   return attach(attachedHwnd);
}

int ProcessDevice::detach_temp()
{
    bool success = AttachThreadInput(GetCurrentThreadId(),
    GetWindowThreadProcessId(attachedHwnd, NULL), false);

    return success;
}

HWND ProcessDevice::getAttachedHwnd()
{
    if( attached )
        return attachedHwnd;
    else
        return NULL;
}

HANDLE ProcessDevice::openProcess(int proc)
{
	#ifdef PROCESS_PROTECT_MODE
		if( proc == (int)GetCurrentProcessId() )
			return NULL;
	#endif

    HANDLE process = OpenProcess(PROCESS_VM_OPERATION | PROCESS_VM_READ |
        PROCESS_VM_WRITE, false, proc);

    return process;
}

void ProcessDevice::closeProcess(HANDLE &proc)
{
    CloseHandle(proc);
    proc = NULL;
}

DWORD ProcessDevice::findProcess(std::string name, int &err)
{
    err = 0;
    HWND handle = findWindow(name, "", err);

    if( handle == 0 ) {
        err = GetLastError();
        return DWORD(0);
    }

    DWORD id;
    GetWindowThreadProcessId(handle, &id);
    return id;
}

DWORD ProcessDevice::findProcessByExe(std::string inname, int &err)
{
    err = 0;
    std::string procname = "";
    DWORD processes[8192];
    DWORD cb;

    int success = EnumProcesses(processes, sizeof(processes), &cb );
    if( !success ) {
        err = GetLastError();
        return DWORD(0);
    }

    DWORD proccount = cb / sizeof(DWORD);
    for(unsigned int i = 0; i < proccount; i++)
    {
        if( processes[i] == 0 ) // skip invalid entries
            continue;

        TCHAR szProcName[MAX_PATH] = TEXT("");
        HANDLE handle = OpenProcess(PROCESS_QUERY_INFORMATION
            | PROCESS_VM_READ, FALSE, processes[i]);

        if( handle == NULL )
        {
            // skip access denied errors
            if( GetLastError() != ERROR_ACCESS_DENIED ) {
            err = GetLastError();
            return DWORD(0); }
        }

        procname = "";
        if( handle != NULL )
        {
            HMODULE hMod;
            DWORD cb;

            // skip errors here...because system processes are very likely
            // to generate errors
            if( EnumProcessModules(handle, &hMod, sizeof(hMod), &cb) )
                GetModuleBaseName(handle, hMod, szProcName,
                    sizeof(szProcName)/sizeof(TCHAR));

            procname = std::string(szProcName);

            CloseHandle(handle);
        }

        if( wildfind(sztolower(inname), sztolower(procname)) )
        { // match, return this ID
            return processes[i];
        }
    }

    return DWORD(0);
}

DWORD ProcessDevice::findProcessByWindow(HWND handle, int &err)
{
    err = 0;
    DWORD id;
    GetWindowThreadProcessId(handle, &id);

    if( id == 0 )
    {
        err = GetLastError();
        return DWORD(0);
    }

    return id;
}

BOOL CALLBACK findWindowProc(HWND hwnd, LPARAM lparam)
{
    EnumWindowPair *winpair = (EnumWindowPair *)lparam;
    char namestring[2048];
    GetWindowText(hwnd, (char *)&namestring, 2047);

    int match = wildfind(sztolower(winpair->windowname),
        sztolower((char*)&namestring));

    if( match )
    {
        // Looking for just the window, not a specific classname
        if( winpair->classname.compare("") == 0 )
        {
			// Ensure that this isn't a window preview/overlay
			char tmpBuf[256];
			GetClassName(hwnd, (char*)&tmpBuf, 256);

			if( !strcmp((char*)&tmpBuf, windowThumbnailClassName))
			return true;

            winpair->hwnd = hwnd;
            return false;
        }
        else
        {
			// Check if this window is valid itself
			char tmpBuf[256];
			GetClassName(hwnd, (char*)&tmpBuf, 256);
			if( strcmp(tmpBuf, winpair->classname.c_str()) == 0 )
			{
				// We have a match
				winpair->hwnd = hwnd;
				return false;
			}

			// If not, scan it's children
            HWND controlHwnd = FindWindowEx(hwnd, NULL,
                winpair->classname.c_str(), NULL);

            if( controlHwnd == NULL )
                return true;

            // We have a match
            winpair->hwnd = controlHwnd;
            return false;
        }
    }
    else
        return true;
}

HWND ProcessDevice::getWindowParent(HWND child)
{
	return GetParent(child);
}

HWND ProcessDevice::findWindow(std::string name, std::string classname,
    int &err)
{
    err = 0;
    EnumWindowPair searchpair;
    searchpair.windowname = name;
    searchpair.classname = classname;
    searchpair.hwnd = 0;

    EnumWindows(findWindowProc, (LPARAM)&searchpair);

    if( searchpair.hwnd == 0 )
        err = GetLastError();

    return searchpair.hwnd;
}


BOOL CALLBACK findWindowListProc(HWND hwnd, LPARAM lParam)
{
    EnumWindowListPair *winpair = (EnumWindowListPair *)lParam;

    char *namestring = new char[2048];
	if( namestring == NULL )
		allocationError();
    GetWindowText(hwnd, namestring, 2047);

    int match = wildfind(sztolower(winpair->windowname), sztolower(namestring));

	delete [] namestring;

    if( match )
    {
        // Looking for just the window, not a specific classname
        if( winpair->classname.compare("") == 0 )
        {
			// Ensure that this isn't a window preview/overlay
			char tmpBuf[256];
			GetClassName(hwnd, (char*)&tmpBuf, 256);

			if( !strcmp((char*)&tmpBuf, windowThumbnailClassName) )
			return true;

            winpair->hwndVec.push_back(hwnd);
        }
        else
        {
			// Check if this window is valid itself
			char tmpBuf[256];
			GetClassName(hwnd, (char*)&tmpBuf, 256);
			if( strcmp(tmpBuf, winpair->classname.c_str()) == 0 )
			{
				// We have a match
				winpair->hwndVec.push_back(hwnd);
			}

            HWND controlHwnd = FindWindowEx(hwnd, NULL,
                winpair->classname.c_str(), NULL);

            if( controlHwnd != NULL )
                winpair->hwndVec.push_back(controlHwnd);
        }
    }

    return true;
}

void ProcessDevice::findWindowList(std::string name, std::string classname,
    std::vector<HWND> &list, int &err)
{
    err = 0;
    EnumWindowListPair searchpair;
    searchpair.windowname = name;
    searchpair.classname = classname;

    BOOL result = EnumWindows(findWindowListProc, (LPARAM)&searchpair);

    if( searchpair.hwndVec.size() == 0 || result == 0 )
        err = GetLastError();

    list = searchpair.hwndVec;
}

HWND ProcessDevice::getHwnd()
{
    return ::getAppHwnd();
}

std::string ProcessDevice::getWindowName(HWND hwnd, int &err)
{
    err = 0;
    char buf[2048];
    if( !GetWindowText(hwnd, buf, 2047) )
        err = GetLastError();

    return buf;
}

std::string ProcessDevice::getWindowClassName(HWND hwnd, int &err)
{
	err = 0;
	char buf[2048];
	if( !GetClassName(hwnd, buf, 2047) )
		err = GetLastError();
	return buf;
}

void ProcessDevice::setWindowName(HWND hwnd, std::string &name, int &err)
{
    err = 0;

    //if( !SetWindowText(hwnd, name.c_str()) )
    //    err = GetLastError();

	if( hwnd == ::getAppHwnd() )
	{
		if( !SetConsoleTitle(name.c_str()) )
			err = GetLastError();
	}
	else
	{
		if( !SendMessage(hwnd, WM_SETTEXT, (WPARAM)0, (LPARAM)name.c_str()) )
			err = GetLastError();
	}
}

int ProcessDevice::windowValid(HWND hwnd)
{
    return IsWindow(hwnd);
}

HWND ProcessDevice::foregroundWindow()
{
    return GetForegroundWindow();
}

RECT ProcessDevice::windowRect(HWND hwnd)
{
    RECT winrect;
    POINT cpoint;
    cpoint.x = 0; cpoint.y = 0;

    ClientToScreen(hwnd, &cpoint);
    GetClientRect(hwnd, &winrect);

    /* Correct the x,y to the window's position */
    winrect.left = cpoint.x;
    winrect.top = cpoint.y;

    return winrect;
}

WinDC ProcessDevice::openDC(HWND hwnd)
{
    HDC hdc = GetWindowDC(hwnd);
    if( hdc == NULL )
        hwnd = NULL;

    WinDC windc;
    windc.hWnd = hwnd;
    windc.hDc = hdc;

    return windc;
}

void ProcessDevice::closeDC(WinDC &windc)
{
    ReleaseDC(windc.hWnd, windc.hDc);
    windc.hWnd = NULL;
    windc.hDc = NULL;
}

int ProcessDevice::makeColor(unsigned char r, unsigned char g, unsigned char b)
{
    int cr, cb, cg;
    cr = (int)r;
    cg = (int)g;
    cb = (int)b;
    int col = (cr << 16) | (cg << 8) | cb;
    return col;
}

unsigned char ProcessDevice::getR(int col)
{
    return (unsigned char)(col >> 16);
}

unsigned char ProcessDevice::getG(int col)
{
    return (unsigned char)((col >> 8) & 0xFF);
}

unsigned char ProcessDevice::getB(int col)
{
    return (unsigned char)(col & 0xFF);
}

int ProcessDevice::getPixel(WinDC *windc, int x, int y)
{
    HDC hdc;
    hdc = GetDC(NULL);

    POINT point;
    point.x = x; point.y = y;
    ClientToScreen(windc->hWnd, &point);

    COLORREF ref = GetPixel(hdc, point.x, point.y);
    ReleaseDC(windc->hWnd, hdc);

    int col = makeColor(GetRValue(ref), GetGValue(ref), GetBValue(ref));
    return col;
}

void ProcessDevice::setPixel(WinDC *windc, int x, int y, int col)
{
    int r, g, b;
    r = getR(col);
    g = getG(col);
    b = getB(col);
    COLORREF ref = RGB(r, g, b);
    ::SetPixel(windc->hDc, x, y, ref);
}

POINT ProcessDevice::pixelSearch(WinDC *windc, int color, int x1, int y1,
    int x2, int y2, unsigned char accuracy, int step)
{
    POINT retval;
    POINT offset;
    RECT winrect;
    int r1, g1, b1, r2, g2, b2;

    if( step <= 1 ) step = 1;

    // figure out the client offset
    winrect.left = 0; winrect.top = 0;
    offset.x = 0; offset.y = 0;
    GetWindowRect(windc->hWnd, &winrect );
    ClientToScreen(windc->hWnd, &offset);
    offset.x = offset.x - winrect.left;
    offset.y = offset.y - winrect.top;

    retval.x = -1; retval.y = -1; // assume it's not going to be found.
    r1 = getR(color); g1 = getG(color); b1 = getB(color);

    COLORREF ref;

    for(int i = y1; i <= y2; i += step)
    {
        for(int v = x1; v <= x2; v += step)
        {
            ref = GetPixel(windc->hDc, offset.x + v, offset.y + i);
            r2 = GetRValue(ref); g2 = GetGValue(ref); b2 = GetBValue(ref);

            if( abs(r2 - r1) <= accuracy &&
                abs(g2 - g1) <= accuracy &&
                abs(b2 - b1) <= accuracy )
            {
                retval.x = v; retval.y = i;
                return retval;
            }
        }
    }

    return retval;
}

void ProcessDevice::drawLine(HDC hdc, int x1, int y1, int x2, int y2,
    int col, int thickness)
{
    HPEN pen;
    LOGBRUSH lb;
    lb.lbStyle = BS_SOLID;
    lb.lbColor = RGB( getR(col), getG(col), getB(col) );
    lb.lbHatch = 0;
    pen = ExtCreatePen(PS_GEOMETRIC | PS_SOLID, thickness, &lb, 0, NULL);

    SelectObject(hdc, pen);

    MoveToEx(hdc, x1, y1, (LPPOINT) NULL);
    LineTo(hdc, x2, y2);

    DeleteObject(pen);
}

void ProcessDevice::drawRect(HDC hdc, int x1, int y1, int x2, int y2,
    int col, int thickness)
{
    HPEN pen;
    LOGBRUSH lb;
    lb.lbStyle = BS_SOLID;
    lb.lbColor = RGB( getR(col), getG(col), getB(col) );
    lb.lbHatch = 0;
    pen = ExtCreatePen(PS_GEOMETRIC | PS_SOLID, thickness, &lb, 0, NULL);

    SelectObject(hdc, pen);

    MoveToEx(hdc, x1, y1, (LPPOINT) NULL);
    LineTo(hdc, x1, y2);

    MoveToEx(hdc, x1, y2, (LPPOINT) NULL);
    LineTo(hdc, x2, y2);

    MoveToEx(hdc, x2, y2, (LPPOINT) NULL);
    LineTo(hdc, x2, y1);

    MoveToEx(hdc, x2, y1, (LPPOINT) NULL);
    LineTo(hdc, x1, y1);

    DeleteObject(pen);
}

void ProcessDevice::saveScreenshot(HWND hwnd, const char *filename)
{
    HDC hdc;
    RECT rect;
    HDC tmpdc;
    HBITMAP bmp;

    if( hwnd == 0 )
        hwnd = GetDesktopWindow();

    /* Get the window rect; change right/bottom to width/height */
    GetWindowRect(hwnd, &rect);
    rect.right = rect.right - rect.left;
    rect.bottom = rect.bottom - rect.top;

    hdc = GetDC(NULL /*hwnd*/);
    tmpdc = CreateCompatibleDC(hdc);
    bmp = CreateCompatibleBitmap(hdc, rect.right, rect.bottom);
    SelectObject(tmpdc, bmp);
    BitBlt(tmpdc, 0, 0, rect.right, rect.bottom, hdc, rect.left, rect.top, SRCCOPY);
    //ReleaseDC(hwnd, hdc);


    BITMAPINFO bmi;
    ZeroMemory(&bmi, sizeof(bmi));

    bmi.bmiHeader.biSize = sizeof(BITMAPINFOHEADER);
    bmi.bmiHeader.biWidth = rect.right;
    bmi.bmiHeader.biHeight = rect.bottom;
    bmi.bmiHeader.biBitCount = 24;
    bmi.bmiHeader.biPlanes = 1;
    bmi.bmiHeader.biCompression = BI_RGB;
    bmi.bmiHeader.biSizeImage = 32 * rect.right * rect.bottom / 8;

    BYTE *pbBits = new BYTE[bmi.bmiHeader.biSizeImage];
	if( pbBits == NULL )
		allocationError();

    GetDIBits(tmpdc, bmp, 0, bmi.bmiHeader.biHeight,
        pbBits, &bmi, DIB_RGB_COLORS);

    BITMAPFILEHEADER bfh;
    bfh.bfType = ('M' << 8) + 'B'; // BM header for bitmaps, always
    bfh.bfSize = sizeof(BITMAPFILEHEADER) + bmi.bmiHeader.biSizeImage +
        sizeof(BITMAPINFOHEADER);
    bfh.bfReserved1 = 0;
    bfh.bfReserved2 = 0;
    bfh.bfOffBits = sizeof(BITMAPFILEHEADER) + sizeof(BITMAPINFOHEADER);

    HANDLE file = CreateFile(filename, GENERIC_WRITE, 0, 0,
        OPEN_ALWAYS, 0, 0);
    DWORD dwWritten;
    WriteFile(file, &bfh, sizeof(bfh), &dwWritten, NULL);
    WriteFile(file, &bmi.bmiHeader, sizeof(BITMAPINFOHEADER), &dwWritten, NULL);
    WriteFile(file, pbBits, bmi.bmiHeader.biSizeImage, &dwWritten, NULL);
    CloseHandle(file);

   //SelectObject(tmpdc, );
    DeleteDC(tmpdc);
    ReleaseDC(hwnd, hdc);
    DeleteObject(bmp);
    delete []pbBits;
}

std::string ProcessDevice::getClipboard()
{
	HGLOBAL cbH = NULL;
	std::string val;

	if( IsClipboardFormatAvailable(CF_TEXT) )
	{
		OpenClipboard(getAppHwnd());
		cbH = GetClipboardData(CF_TEXT);

		if( !cbH )
			return "";

		PSTR pClip = (CHAR *)GlobalLock(cbH);
		if( !pClip ) {
			CloseClipboard();
			return "";
		}

		val = pClip;

		GlobalUnlock(cbH);
		CloseClipboard();
	}

	return val;
}

void ProcessDevice::setClipboard(std::string &str)
{
	HGLOBAL cbH = NULL;

	cbH = GlobalAlloc(GMEM_MOVEABLE, str.size() + 1);
	if( !cbH )
		return;

	PSTR pClip = (CHAR *)GlobalLock(cbH);
	if( !pClip ) {
		GlobalFree(cbH);
		return;
	}

	strcpy(pClip, str.c_str());

	GlobalUnlock(pClip);

	OpenClipboard(getAppHwnd());
	EmptyClipboard();
	SetClipboardData(CF_TEXT, pClip);
	CloseClipboard();
}

void ProcessDevice::showWindow(HWND window, int cmd)
{
    ShowWindow(window, cmd);

    if( cmd == SW_SHOW || cmd == SW_SHOWNORMAL || cmd == SW_RESTORE )
        SetForegroundWindow(window);
}

bool ProcessDevice::procDataCompare(const unsigned char *data,
    const unsigned char *bmask, const char *szMask)
{
    for(; *szMask; ++szMask, ++data, ++bmask)
    {
        if( *szMask == 'x' && *data!= *bmask )
            return false;
    }
    return (*szMask) == 0;
}


unsigned long ProcessDevice::findPatternInProcess(HANDLE proc, unsigned char *bmask,
    char *szMask, unsigned long address = 0, unsigned long len = 4)
{
    unsigned long retval = 0;
    unsigned long dataLen = strlen(szMask);
    SIZE_T bytesRead;
	unsigned long buffStart = 0;//address;
	unsigned int bufferLen = dataLen * 50;
	if( bufferLen < 1024 )
		bufferLen = 1024;
	int addLen = bufferLen % dataLen;
	if( addLen > 0 )
		bufferLen = bufferLen + addLen;

	unsigned int curBufferLen = bufferLen; // We may not always read the full length

	unsigned char *buffer = new unsigned char[bufferLen + 1];
	if( buffer == NULL )
		allocationError();
	//memset(buffer, 0, bufferLen);

    for(unsigned long i = 0; i < len; i++)
    {
        unsigned int curAddr = address + i;
        /*bool success = ReadProcessMemory(proc, (LPCVOID)curAddr, data,
            dataLen, &bytesRead);*/
		if( (curAddr - buffStart + dataLen) >= curBufferLen ) // We need to advance the buffer
		{
			memset(buffer, 0, bufferLen);
			bool success = ReadProcessMemory(proc, (LPCVOID)curAddr, buffer,
				bufferLen, &bytesRead);

			buffStart = curAddr;

			if( !success && bytesRead == 0 )
				continue;

			#ifdef DISPLAY_DEBUG_MESSAGES
			if( !success && bytesRead <  bufferLen )
			{
					Logger::instance()->add("DEBUG: findPatternInProcess() did not read full length. Got %d bytes, expected %d.",
						(int)bytesRead, bufferLen);
			}
			#endif

			curBufferLen = bytesRead;
		}
		//memcpy(data, (char *)&buffer[curAddr - buffStart],dataLen);

        if( procDataCompare((const unsigned char *)&buffer[curAddr - buffStart], bmask, szMask) )
        {
            retval = address + i;
            break;
        }
    }

//    delete [] data;
	delete [] buffer;
    return retval;
}

unsigned long ProcessDevice::getModuleAddress(DWORD proc, const char *modname)
{
	char modname_lower[256];
	cstrtolower(modname_lower, modname, 256);

	HANDLE snapshot = CreateToolhelp32Snapshot(TH32CS_SNAPMODULE, proc);

	if( snapshot == INVALID_HANDLE_VALUE )
	{
		Logger::instance()->add("Invalid Toolhelp32 snapshot handle returned.");
		return 0;
	}

	MODULEENTRY32 mod;
	mod.dwSize = sizeof(MODULEENTRY32);

	if( Module32First(snapshot, &mod) )
	{
		char modname_snap[256];
		cstrtolower(modname_snap, mod.szModule, 256);

		if( strcmp(modname_snap, modname_lower) == 0 )
			return (unsigned long)mod.modBaseAddr;

		while( Module32Next(snapshot, &mod) )
		{
			//char modname_snap[256];
			cstrtolower(modname_snap, mod.szModule, 256);

			if( strcmp(modname_snap, modname_lower) == 0 )
				return (unsigned long)mod.modBaseAddr;
		}

		/* Failed to find the module */
		return 0;
	}
	else
	{
		/* Failed to read any module info */
		Logger::instance()->add("Failure to read module info.");
		return 0;
	}
}


void ProcessDevice::flashWindow(HWND window, int count)
{
	FLASHWINFO fwi;

	if( window )
		fwi.hwnd = window;
	else
		fwi.hwnd = ::getAppHwnd();

	fwi.dwFlags = FLASHW_ALL;
	if( count < 0 )
		fwi.dwFlags = FLASHW_STOP;
	else if( count == 0 )
		fwi.dwFlags |= FLASHW_TIMERNOFG;
	else
	{
		fwi.dwFlags |= FLASHW_TIMERNOFG;
		fwi.uCount = count;
	}

	fwi.cbSize = sizeof(FLASHWINFO);
	fwi.dwTimeout = 0;

	FlashWindowEx(&fwi);
}
